import time,os,sys,logging
from datetime import datetime
import pandas as pd
from xtquant import xtdata


class WarmupManager:
    """
    预热数据管理器
    负责为各个指标提供历史数据预热功能
    """

    def __init__(self, xtdata, stock_codes, stock_type, log_callback=None):
        self.xtdata = xtdata
        self.stock_codes = stock_codes
        self.stock_type = stock_type
        self.log_callback = log_callback

        # 配置常量
        self.VOLUME_SCALE_THRESHOLD = 100000  # 成交量缩放阈值
        self.VOLUME_SCALE_FACTOR = 100  # 成交量缩放因子
        self.MACD_DOWNLOAD_DELAY = 0.3  # MACD数据下载延迟（秒）
        self.NINE_TURN_DOWNLOAD_DELAY = 1.0  # 九转数据下载延迟（秒）

        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()
        
        # 创建成交量指标实例
        from 主程.指标.成交量 import solution
        self.volume_solution = solution()

    def log(self, msg):
        """统一的日志记录方法"""
        print(msg)
        if self.log_callback:
            self.log_callback(msg)

    def warmup_all_indicators(self, nine_turn, rsi_calculator, volume_solution, macd_monitor=None, enable_warmup=True, trade_date=None):
        """
        预热所有指标的历史数据

        Args:
            nine_turn: 九转指标计算器
            rsi_calculator: RSI计算器
            volume_solution: 成交量指标计算器
            macd_monitor: MACD监控器
            enable_warmup: 是否启用预热
            trade_date: 回测日期（格式：YYYYMMDD），如果为None则使用当前日期
        """
        self.log("🔥 开始预热指标数据...")
        
        # 确定预热日期
        if trade_date is None:
            warmup_date = datetime.now().strftime("%Y%m%d")
            self.log(f"📅 使用当前日期进行预热: {warmup_date}")
        else:
            warmup_date = trade_date
            self.log(f"📅 使用回测日期进行预热: {warmup_date}")

        # 1. 九转指标预热
        self.warmup_nine_turn(nine_turn, warmup_date)

        # 2. RSI和成交量指标预热
        if enable_warmup:
            self.warmup_rsi_and_volume(rsi_calculator, volume_solution, warmup_date)

        # 3. MACD指标预热
        if macd_monitor and enable_warmup:
            self.warmup_macd(macd_monitor, warmup_date)

        self.log("✅ 指标数据预热完成")

    def warmup_nine_turn(self, nine_turn, warmup_date=None):
        """
        九转指标预热：专门为513090.SH预热九转数据

        Args:
            nine_turn: 九转指标计算器
            warmup_date: 预热日期（格式：YYYYMMDD），如果为None则使用当前日期
        """
        try:
            _prev_verbose = getattr(nine_turn, 'verbose', True)
            try:
                nine_turn.verbose = False  # 预热阶段关闭九转日志，避免刷屏
            except Exception:
                pass

            code = "513090.SH"
            # 使用传入的预热日期，如果没有则使用当前日期
            today = warmup_date if warmup_date else datetime.now().strftime("%Y%m%d")

            # 下载当日tick数据
            try:
                self.xtdata.download_history_data(stock_code=code, period="tick", start_time=today, end_time=today)
                time.sleep(self.NINE_TURN_DOWNLOAD_DELAY)
            except Exception:
                pass

            # 获取并处理历史数据
            data = self.xtdata.get_local_data(stock_list=[code], period='tick', start_time=today, end_time=today)
            df = data.get(code) if data else None

            if df is not None and not df.empty:
                time_col = next((c for c in df.columns if 'time' in c.lower()), None)
                price_col = next((c for c in ['price', 'lastPrice', 'close', 'last', 'current'] if c in df.columns),
                                 None)

                if time_col and price_col:
                    # 时间格式统一处理
                    if str(df[time_col].dtype) in ('int64', 'float64'):
                        df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
                    else:
                        df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')

                    df = df[df['__time__'].notna()].sort_values('__time__')
                    hist = pd.DataFrame({'price': df[price_col].astype(float).values}, index=df['__time__'])

                    # 预热九转指标
                    nine_turn.warmup_from_history(code, hist)
                    self.log("✅ 九转预热完成：513090.SH")
                else:
                    self.log("❌ 九转预热失败：无法识别时间或价格列")
            else:
                self.log("❌ 九转预热失败或无数据：513090.SH")

        except Exception as e:
            self.log(f"❌ 九转预热异常：513090.SH - {e}")
        finally:
            try:
                nine_turn.verbose = _prev_verbose  # 恢复原始日志设置
            except Exception:
                pass

    def warmup_rsi_and_volume(self, rsi_calculator, volume_solution, warmup_date=None):
        """
        RSI和成交量指标预热：为所有股票填充当日分钟历史数据

        Args:
            rsi_calculator: RSI计算器
            volume_solution: 成交量指标计算器
            warmup_date: 预热日期（格式：YYYYMMDD），如果为None则使用当前日期
        """
        try:
            # 使用传入的预热日期，如果没有则使用当前日期
            today = warmup_date if warmup_date else datetime.now().strftime("%Y%m%d")
            # 在回测模式下，使用预热日期；在实时模式下，使用当前时间
            if warmup_date:
                now_floor = pd.to_datetime(f"{warmup_date} 15:00")  # 回测模式：使用回测日期的收盘时间
            else:
                now_floor = pd.to_datetime(datetime.now().strftime('%Y-%m-%d %H:%M'))  # 实时模式：使用当前时间

            for code in self.stock_codes:
                try:
                    # 下载当日分钟数据（更适合RSI计算）
                    self._download_minute_data(code, today)

                    # 获取并处理分钟数据
                    df = self._get_and_process_minute_data(code, today)
                    if df is None:
                        continue

                    # RSI预热（新的独立RSI计算器会自动处理，这里只是预下载数据）
                    self._pre_download_rsi_data(code, today)

                    # 成交量预热（仍使用tick数据，因为成交量需要实时性）
                    self._warmup_volume_for_stock_tick(code, today, volume_solution)

                except Exception as e:
                    self.log(f"⚠️ {code} 预热异常: {e}")

        except Exception as e:
            self.log(f"❌ RSI和成交量预热异常: {e}")

    def warmup_macd(self, macd_monitor, warmup_date=None):
        """
        MACD指标预热：为所有股票下载历史数据并初始化MACD计算器

        Args:
            macd_monitor: MACD监控器实例
            warmup_date: 预热日期（格式：YYYYMMDD），如果为None则使用当前日期
        """
        try:
            self.log("🔥 开始MACD指标预热...")
            
            # 使用传入的预热日期，如果没有则使用当前日期
            today = warmup_date if warmup_date else datetime.now().strftime("%Y%m%d")
            
            for code in self.stock_codes:
                if code in ("513130.SH", "513060.SH"):
                    continue  # 跳过波动区间股票
                    
                try:
                    self.log(f"🔄 {code} MACD预热中...")
                    
                    # 下载历史数据（5天数据用于MACD预热）
                    try:
                        # 在回测模式下，基于预热日期计算历史日期；在实时模式下，基于当前日期计算
                        if warmup_date:
                            # 回测模式：基于回测日期计算历史日期
                            warmup_dt = datetime.strptime(warmup_date, "%Y%m%d")
                            start_date = (warmup_dt - pd.Timedelta(days=5)).strftime("%Y%m%d")
                        else:
                            # 实时模式：基于当前日期计算历史日期
                            start_date = (datetime.now() - pd.Timedelta(days=5)).strftime("%Y%m%d")
                            
                        self.xtdata.download_history_data(
                            stock_code=code, 
                            period="1m",  # 使用1分钟数据，更适合MACD计算
                            start_time=start_date, 
                            end_time=today
                        )
                        time.sleep(self.MACD_DOWNLOAD_DELAY)
                    except Exception as e:
                        self.log(f"⚠️ {code} MACD历史数据下载失败: {e}")
                        continue
                    
                    # 获取历史数据
                    try:
                        data = self.xtdata.get_local_data(
                            stock_list=[code], 
                            period='1m', 
                            start_time=start_date, 
                            end_time=today
                        )
                        df = data.get(code) if data else None
                        
                        if df is not None and not df.empty:
                            # 确保有足够的数据点进行MACD计算
                            if len(df) >= 50:  # MACD需要至少50个数据点
                                self.log(f"✅ {code} MACD预热完成，数据点数: {len(df)}")
                            else:
                                self.log(f"⚠️ {code} MACD数据不足，点数: {len(df)}")
                        else:
                            self.log(f"⚠️ {code} MACD无历史数据")
                            
                    except Exception as e:
                        self.log(f"⚠️ {code} MACD历史数据获取失败: {e}")
                        
                except Exception as e:
                    self.log(f"⚠️ {code} MACD预热异常: {e}")
                    
            self.log("✅ MACD指标预热完成")
            
        except Exception as e:
            self.log(f"❌ MACD指标预热异常: {e}")

    def _download_tick_data(self, code: str, today: str):
        """下载指定股票的当日tick数据"""
        try:
            self.xtdata.download_history_data(stock_code=code, period="tick", start_time=today, end_time=today)
            time.sleep(self.MACD_DOWNLOAD_DELAY)
        except Exception:
            pass
    
    def _download_minute_data(self, code: str, today: str):
        """下载指定股票的当日分钟数据（用于RSI预热）"""
        try:
            self.xtdata.download_history_data(stock_code=code, period="1m", start_time=today, end_time=today)
            time.sleep(self.MACD_DOWNLOAD_DELAY)
        except Exception:
            pass

    def _get_and_process_tick_data(self, code: str, today: str):
        """获取并处理tick数据，返回处理后的DataFrame"""
        try:
            data = self.xtdata.get_local_data(stock_list=[code], period='tick', start_time=today, end_time=today)
            df = data.get(code) if data else None

            if df is None or df.empty:
                return None

            # 识别时间、价格、成交量列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)
            vol_col = next((c for c in ['volume', 'vol', '成交量'] if c in df.columns), None)

            if not time_col or not price_col:
                return None

            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')

            df = df[df['__time__'].notna()].sort_values('__time__')

            # 添加列名到DataFrame以便后续使用
            df._price_col = price_col
            df._vol_col = vol_col

            return df if not df.empty else None

        except Exception:
            return None
    
    def _get_and_process_minute_data(self, code: str, today: str):
        """获取并处理分钟数据，返回处理后的DataFrame（用于RSI预热）"""
        try:
            data = self.xtdata.get_local_data(stock_list=[code], period='1m', start_time=today, end_time=today)
            df = data.get(code) if data else None

            if df is None or df.empty:
                return None

            # 识别时间、价格列
            time_col = next((c for c in df.columns if 'time' in c.lower()), None)
            price_col = next((c for c in ['close', 'lastPrice', 'price', 'last', 'current'] if c in df.columns), None)

            if not time_col or not price_col:
                return None

            # 时间格式统一处理
            if str(df[time_col].dtype) in ('int64', 'float64'):
                df['__time__'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                try:
                    df['__time__'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
                except Exception:
                    df['__time__'] = pd.to_datetime(df[time_col], errors='coerce')

            df = df[df['__time__'].notna()].sort_values('__time__')

            # 添加列名到DataFrame以便后续使用
            df._price_col = price_col

            return df if not df.empty else None

        except Exception:
            return None

        # 旧的RSI预热方法已移除，新的独立RSI计算器会自动处理
    
    def _pre_download_rsi_data(self, code: str, today: str):
        """预下载RSI数据（新的独立RSI计算器会自动处理）"""
        try:
            # 新的独立RSI计算器会自动下载和处理数据
            # 这里只是预下载，让数据在本地可用
            self._download_minute_data(code, today)
            self.log(f"✅ {code} RSI数据预下载完成")
        except Exception as e:
            self.log(f"⚠️ {code} RSI数据预下载失败: {e}")

    def _warmup_volume_for_stock(self, code: str, df: pd.DataFrame, minute_bucket: pd.Series, now_floor: pd.Timestamp,
                                 volume_solution):
        """为单个股票预热成交量数据"""
        try:
            vol_col = df._vol_col
            if not vol_col:
                return

            # 计算每分钟的增量成交量
            minute_increments = self._calculate_minute_volume_increments(df, vol_col, minute_bucket, now_floor)

            if not minute_increments:
                return

            # 转换为Series格式
            minute_vol = pd.Series(dict(minute_increments))
            minute_vol = minute_vol[minute_vol.index < now_floor]

            # 调试信息
            if len(minute_vol) > 0:
                self.log(
                    f"🔍 {code} 预热成交量数据: 范围[{minute_vol.min():.0f}, {minute_vol.max():.0f}], 平均值{minute_vol.mean():.0f}")

            # 筛选交易时段内的数据
            minute_vol = self._filter_trading_session_data(minute_vol)

            # 数据缩放处理
            minute_vol = self._scale_volume_data(code, minute_vol)

            # 预热成交量指标
            if len(minute_vol) > 0:
                try:
                    volume_solution.warmup_with_volumes(code, minute_vol.tolist())
                    self.log(f"✅ {code} 成交量预热完成，数据点数: {len(minute_vol)}")
                except Exception as e:
                    self.log(f"⚠️ {code} 成交量预热失败: {e}")

        except Exception as e:
            self.log(f"⚠️ {code} 成交量预热异常: {e}")
    
    def _warmup_volume_for_stock_tick(self, code: str, today: str, volume_solution):
        """为单个股票预热成交量数据（基于tick数据）"""
        try:
            # 下载当日tick数据
            self._download_tick_data(code, today)
            
            # 获取并处理tick数据
            df = self._get_and_process_tick_data(code, today)
            if df is None:
                return

            vol_col = df._vol_col
            if not vol_col:
                return

            # 计算分钟聚合
            minute_bucket = df['__time__'].dt.floor('min')
            now_floor = pd.to_datetime(datetime.now().strftime('%Y-%m-%d %H:%M'))

            # 计算每分钟的增量成交量
            minute_increments = self._calculate_minute_volume_increments(df, vol_col, minute_bucket, now_floor)

            if not minute_increments:
                return

            # 转换为Series格式
            minute_vol = pd.Series(dict(minute_increments))
            minute_vol = minute_vol[minute_vol.index < now_floor]

            # 筛选交易时段内的数据
            minute_vol = self._filter_trading_session_data(minute_vol)

            # 数据缩放处理
            minute_vol = self._scale_volume_data(code, minute_vol)

            # 预热成交量指标
            if len(minute_vol) > 0:
                try:
                    volume_solution.warmup_with_volumes(code, minute_vol.tolist())
                    self.log(f"✅ {code} 成交量tick数据预热完成，数据点数: {len(minute_vol)}")
                except Exception as e:
                    self.log(f"⚠️ {code} 成交量tick数据预热失败: {e}")

        except Exception as e:
            self.log(f"⚠️ {code} 成交量tick数据预热异常: {e}")

    def _calculate_minute_volume_increments(self, df: pd.DataFrame, vol_col: str, minute_bucket: pd.Series,
                                            now_floor: pd.Timestamp):
        """计算每分钟的成交量增量"""
        minute_increments = []

        try:
            for minute, group in df.groupby(minute_bucket):
                if minute < now_floor:  # 排除当前未完成分钟
                    # 计算这一分钟内成交量的增量
                    minute_volumes = group[vol_col].tolist()
                    minute_volumes_numeric = [float(vol) for vol in minute_volumes if pd.notna(vol)]

                    if len(minute_volumes_numeric) > 1:
                        # 如果这一分钟内有多个tick，计算增量
                        increment = minute_volumes_numeric[-1] - minute_volumes_numeric[0]
                    else:
                        # 如果只有一个tick，增量就是该值
                        increment = minute_volumes_numeric[0] if minute_volumes_numeric else 0

                    # 确保增量不为负数
                    if increment >= 0:
                        minute_increments.append((minute, increment))

        except Exception:
            pass

        return minute_increments

    def _filter_trading_session_data(self, minute_vol: pd.Series) -> pd.Series:
        """筛选交易时段内的数据"""
        try:
            mask = [((self.MORNING_START <= ts.time() <= self.MORNING_END) or
                     (self.AFTERNOON_START <= ts.time() <= self.AFTERNOON_END)) for ts in minute_vol.index]
            if any(mask):
                return minute_vol[mask]
        except Exception:
            pass
        return minute_vol

    def _scale_volume_data(self, code: str, minute_vol: pd.Series) -> pd.Series:
        """对成交量数据进行缩放处理"""
        if len(minute_vol) > 0:
            # 如果数据异常大，可能是单位问题，进行缩放
            if minute_vol.max() > self.VOLUME_SCALE_THRESHOLD:
                self.log(f"⚠️ {code} 预热成交量数据异常大，进行缩放处理")
                minute_vol = minute_vol / self.VOLUME_SCALE_FACTOR
        return minute_vol


class DataProcessor:
    """
    数据处理器
    负责处理分钟聚合、交易时段判断等通用功能
    """

    def __init__(self):
        # 交易时间段常量
        self.MORNING_START = datetime.strptime("09:30", "%H:%M").time()
        self.MORNING_END = datetime.strptime("11:30", "%H:%M").time()
        self.AFTERNOON_START = datetime.strptime("13:00", "%H:%M").time()
        self.AFTERNOON_END = datetime.strptime("15:00", "%H:%M").time()

    def to_minute_bucket(self, ts) -> str:
        """
        将时间戳转换为分钟聚合格式

        Args:
            ts: 时间戳，支持多种格式

        Returns:
            str: 分钟聚合字符串，格式为 'YYYY-MM-DD HH:MM'
        """
        try:
            if ts is None:
                return None
            if isinstance(ts, (int, float)):
                t = float(ts)
                if t > 1e12:
                    dt = datetime.fromtimestamp(t / 1000)
                else:
                    dt = datetime.fromtimestamp(t)
                return dt.strftime('%Y-%m-%d %H:%M')
            s = str(ts)
            if s.isdigit() and len(s) == 14:
                dt = datetime.strptime(s, '%Y%m%d%H%M%S')
                return dt.strftime('%Y-%m-%d %H:%M')
            try:
                dt = pd.to_datetime(s, errors='coerce')
                if pd.isna(dt):
                    return None
                if getattr(dt, 'tzinfo', None) is not None:
                    dt = dt.tz_convert('Asia/Shanghai').tz_localize(None)
                return dt.strftime('%Y-%m-%d %H:%M')
            except Exception:
                return None
        except Exception:
            return None

    def is_session_minute(self, minute_str: str) -> bool:
        """
        判断是否为交易时段内的分钟

        Args:
            minute_str: 分钟字符串，格式为 'YYYY-MM-DD HH:MM'

        Returns:
            bool: 是否为交易时段内的分钟
        """
        try:
            if not minute_str:
                return False
            dt = datetime.strptime(minute_str, '%Y-%m-%d %H:%M')
            t = dt.time()
            return (self.MORNING_START <= t <= self.MORNING_END) or (
                        self.AFTERNOON_START <= t <= self.AFTERNOON_END)
        except Exception:
            return False

    def is_trading_time(self) -> bool:
        """
        检查当前是否在交易时间段内

        Returns:
            bool: 是否在交易时间段内
        """
        try:
            now = datetime.now()
            current_time = now.time()

            # 检查是否在交易时间段内
            is_morning_trading = self.MORNING_START <= current_time <= self.MORNING_END
            is_afternoon_trading = self.AFTERNOON_START <= current_time <= self.AFTERNOON_END

            return is_morning_trading or is_afternoon_trading

        except Exception:
            return False


# 创建全局的data_processor实例
data_processor = DataProcessor()